# coding=utf-8
from __future__ import absolute_import, division, print_function, unicode_literals

import io
import ipaddress
import os
import re

import servicediscovery
from bytedmysql import sqlalchemy_init
from sqlacodegen.codegen import CodeGenerator
from sqlalchemy import create_engine, inspect
from sqlalchemy.schema import MetaData

from apitest.engine.config import main_config
from apitest.engine.const.enums import IP
from apitest.engine.const.model import FilterRule, RuleType

ROOT_PATH = os.getcwd()


class GenMysqlModel:

    def __init__(self, name, out=None, filter_rule: FilterRule = None):
        self.db = main_config.db.get(name)
        self.filter_rule = filter_rule
        self.out = out
        if not self.out:
            self.out = os.path.join(ROOT_PATH, "model", "mysql")
            if not os.path.exists(self.out):
                os.makedirs(self.out)
            self.out = os.path.join(self.out, "{}.py".format(self.db.name))
        else:
            if self.out.split(".")[-1] != "py":
                self.out = os.path.join(self.out, "{}.py".format(self.db.name))

    def get_db_info(self):
        # 优先ipv6协议，其次ipv4
        try:
            db_info = servicediscovery.lookup(self.db.psm, address_family="dual-stack")[0]
        except IndexError:
            try:
                db_info = servicediscovery.lookup(self.db.psm)[0]
            except IndexError:
                raise Exception("get db host error")
        return db_info

    def genMysqlModel(self):
        sqlalchemy_init()
        dbInfo = self.get_db_info()
        host = dbInfo.host
        IPv = ipaddress.ip_address(dbInfo.host).version
        if IPv != IP.IPv6.value and IPv!= IP.IPv4.value:
            raise Exception("IP version is not IPv4 or IPv6")
        if IPv == IP.IPv6.value:
            host = "[{}]".format(host)
        engine = create_engine(
            f"mysql+pymysql://{self.db.user}:{self.db.pwd}@{host}:{str(dbInfo.port)}/?charset=utf8mb4&&disable_metrics=True"
            )

        inspector = inspect(engine)
        available = inspector.get_table_names()
        tables = []
        if self.filter_rule:
            if self.filter_rule.rule_type == RuleType.regex:
                if not isinstance(self.filter_rule.rule_code, str):
                    raise Exception("filter rule must str")
                try:
                    pattern = re.compile(r"{}".format(self.filter_rule.rule_code))
                except re.error:
                    raise Exception("filter code error")
                tables = [tab for tab in available if pattern.match(tab)]
            elif self.filter_rule.rule_type == RuleType.keywords:
                if isinstance(self.filter_rule.rule_code, str):
                    self.filter_rule.rule_code = [self.filter_rule.rule_code]
                for keyword in self.filter_rule.rule_code:
                    for tab in available:
                        if keyword in tab:
                            tables.append(tab)
                tables = list(set(tables))
            else:
                if not isinstance(self.filter_rule.rule_code, list):
                    raise Exception("filter code type must List[str]")
                tables = self.filter_rule.rule_code

        metadata = MetaData(engine)
        if tables:
            metadata.reflect(bind=engine, schema=self.db.name, only=tables)
        else:
            metadata.reflect(bind=engine, schema=self.db.name)

        # Write the generated model code to the specified file or standard output
        outfile = io.open(self.out, 'w', encoding='utf-8')
        generator = CodeGenerator(metadata=metadata)
        generator.render(outfile)
